import os
import numpy as np
from easy_io import write_pkl_file
import global_config


fold_list = np.load('/data_4t/Kaggle/fold_list.npy')
label_dict = np.load('/data_4t/Kaggle/lidc&kaggle/label_dict3.npy').item()

nb_total_scans_dict, nb_total_labels_dict = {}, {}

for i in range(len(fold_list)):
    scans = fold_list[i]
    nb_total_scans_dict[i] = len(scans)
    nb_total_labels_dict[i] = sum(len(label_dict[s]) for s in scans)

os.makedirs(global_config.data_folder, exist_ok=True)
write_pkl_file(os.path.join(global_config.data_folder, 'info.pkl'), (nb_total_scans_dict, nb_total_labels_dict))
